# TODO: 导入必要的库和模块
import pickle
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# TODO: 加载数字数据集（使用sklearn内置的手写数字数据集，与之前识别逻辑匹配）
digits = load_digits()
X = digits.data  # 特征：8x8灰度图的扁平化数组（64维）
y = digits.target  # 标签：0-9的数字

# TODO: 将数据集划分为训练集和测试集（设置random_state确保结果可复现）
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42  # 测试集占比20%，随机种子固定
)

# TODO: 初始化变量以存储最佳准确率，相应的k值和最佳knn模型
best_accuracy = 0.0  # 初始最佳准确率设为0
best_k = 1  # 初始最佳k值设为1
best_knn_model = None  # 初始最佳模型设为None

# TODO: 初始化一个列表以存储每个k值的准确率
k_accuracies = []  # 索引对应k值（k=1对应索引0，k=40对应索引39）

# TODO: 尝试从1到40的k值，训练模型并记录最优结果
k_range = range(1, 41)  # k值范围：1~40
for k in k_range:
    # 初始化KNN分类器（默认欧氏距离，与手写数字识别场景适配）
    knn = KNeighborsClassifier(n_neighbors=k)
    
    # 训练模型（使用训练集）
    knn.fit(X_train, y_train)
    
    # 模型预测（使用测试集）
    y_pred = knn.predict(X_test)
    
    # 计算当前k值的准确率
    current_accuracy = accuracy_score(y_test, y_pred)
    
    # 将当前准确率存入列表
    k_accuracies.append(current_accuracy)
    
    # 对比并更新最佳结果
    if current_accuracy > best_accuracy:
        best_accuracy = current_accuracy
        best_k = k
        best_knn_model = knn  # 保存当前最优模型实例

# TODO: 将最佳KNN模型保存到二进制文件（使用pickle序列化）
with open("best_knn_model.pkl", "wb") as f:  # 'wb'=二进制写入模式
    pickle.dump(best_knn_model, f)
print(f"最佳KNN模型已保存至: best_knn_model.pkl")

# TODO: 打印最佳准确率和相应的k值
print(f"\n===== 最优K值结果 =====")
print(f"最佳k值: {best_k}")
print(f"最佳测试准确率: {best_accuracy:.4f} (即 {best_accuracy*100:.2f}%)")

# 绘制准确率折线图并保存为PDF（含垂直线标注最优k值）
plt.figure(figsize=(10, 6))  # 设置画布大小

# 1. 绘制k值-准确率折线图
plt.plot(k_range, k_accuracies, marker='o', linestyle='-', color='#1f77b4', markersize=4)

# 2. 绘制垂直红线标注最优k值
plt.axvline(x=best_k, color='red', linestyle='--', linewidth=2, label=f'最优k值: {best_k}')

# 3. 在最优k值位置标注准确率（避免文字遮挡，偏移0.5个单位）
plt.text(
    x=best_k + 0.5,  # 文字x坐标（右移0.5）
    y=best_accuracy,  # 文字y坐标（与准确率对齐）
    s=f'k={best_k}\n准确率={best_accuracy:.4f}',  # 标注内容
    fontsize=10,
    bbox=dict(boxstyle='round,pad=0.3', facecolor='yellow', alpha=0.7)  # 黄色背景框
)

# 4. 设置图表标题、坐标轴标签和图例
plt.title('KNN模型不同k值的测试准确率', fontsize=14, pad=20)
plt.xlabel('k值（近邻数量）', fontsize=12)
plt.ylabel('测试准确率', fontsize=12)
plt.legend(loc='lower right', fontsize=10)

# 5. 设置x轴刻度（确保1~40都显示，每5个标注一个）
plt.xticks(range(0, 42, 5), fontsize=10)
# 设置y轴范围（从0.9到1.0，突出准确率差异）
plt.ylim(0.9, 1.0)

# 6. 添加网格线（便于读取数值）
plt.grid(True, linestyle=':', alpha=0.7)

# 7. 保存图表为PDF（bbox_inches避免文字被截断）
plt.savefig('accuracy_plot.pdf', dpi=300, bbox_inches='tight')
plt.close()  # 关闭画布释放资源
print(f"准确率可视化图表已保存至: accuracy_plot.pdf")